#!/usr/bin/env python
# -*- coding:utf-8 -*-
import json

from flask import Response, stream_with_context
from langchain.memory import ConversationBufferMemory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

from controller.pptBaseController import PptBaseController
from models import db
from models.ppt_base import PptBase

from controller.pptChatController import PptChatController
from models.ppt_chat import PptChat
from service.llmService import LLMService
from utils.fetch_json import fetch_json
from utils.loggings import loggings
from utils.response_code import RET


class PptChatService(PptChatController):
    basic_prompt = """你是一个PPT设计师，现在用户想要获取一个PPT的基本框架结构，包含了标题，每个章节，子标题的大纲结构，用户将使用这目录结构来组织自己的PPT。
    请根据用户的需求，生成一个PPT的基本框架结构。结构示例如下：
    {{
        "title": "PPT标题",
        "subtitle": [
        {{
            "title": "章节1",
            "subtitle": [{{
                "title": "子标题1",
                "subtitle": []
            }}, {{
                "title": "子标题2",
                "subtitle": []
            }}]
        }}]
    }}
    请输出一个json格式的结果，结构包裹在```json```标签中。如果有其他的说明性文字需要放在标签之外。可以适当给一些说明文字。
    """
    title_prompt = "请为以下内容起一个标题来概括这段对话，标题要比较简洁明了。以下输入都是在创建一个PPT框架的背景下的。最终直接输出这个标题即可。"

    @classmethod
    def chat_ppt_framework(cls, **kwargs):
        if kwargs.get('ppt_id') is None:
            return {'code': RET.PARAMERR, 'message': 'ppt_id is required', 'data': {}}

        files = kwargs.get('files')
        has_file = files is not None and len(files) > 0

        res_ppt = PptBaseController.get(ppt_id=kwargs.get('ppt_id'))
        if res_ppt['code'] != RET.OK or len(res_ppt['data']) != 1:
            return {'code': RET.DATAERR, 'message': 'PPT not found', 'data': {}}
        ppt = res_ppt['data'][0]
        if ppt['gen_type']:
            if has_file is False:
                if ppt['gen_type'] == 0:
                    return cls.chat_ppt_framework_basic(**kwargs)
                if ppt['gen_type'] == 1:
                    from service.pptChatDocumentService import PptChatDocumentService
                    return PptChatDocumentService.chat_ppt_framework_with_document(**kwargs)
            else:
                from service.pptChatDocumentService import PptChatDocumentService
                return PptChatDocumentService.chat_ppt_framework_with_document(**kwargs)

        else:
            if has_file is False:
                return cls.chat_ppt_framework_basic(**kwargs)
            else:
                from service.pptChatDocumentService import PptChatDocumentService
                return PptChatDocumentService.chat_ppt_framework_with_document(**kwargs)

    @classmethod
    def chat_ppt_framework_basic(cls, **kwargs):
        """
            用户输入较短的文字，生成ppt的文字框架
        """
        from flask import g
        user_id = g.user['user_id']
        content = kwargs.get('content')

        try:
            model = LLMService.get_model()
        except Exception as e:
            return Response({'code': RET.DATAERR, 'message': str(e), 'data': str(e)}, mimetype="text/event-stream",
                            headers={})

        chat_history = cls.get_chat_history(ppt_id=kwargs.get('ppt_id'))
        # First chat
        if len(chat_history['data']) == 0:
            prompt = ChatPromptTemplate.from_messages(cls.get_prompt(cls.basic_prompt, "{user_input}"))

            chain = prompt | model

            response = chain.stream({"user_input": content})

        # Chat history exists
        else:
            history = [('system', cls.basic_prompt)]
            for chat in chat_history['data']:
                history.append(('user', chat['content']))
                history.append(('system', chat['response']))

            prompt = ChatPromptTemplate.from_messages(
                [
                    MessagesPlaceholder(variable_name="history"),
                    ("user", "用户输入：{user_input}"),
                ]
            )
            chain = prompt | model
            response = chain.stream({"user_input": content, "history": history})

        total_response = ''

        def generate():
            nonlocal total_response
            for trunk in response:
                c = trunk.content
                total_response += c

                yield f"event: message\ndata: {json.dumps({'content': c})}\n\n"

            from manage import app
            with app.app_context():

                try:
                    # 加入数据库
                    from utils.generate_id import GenerateID

                    ppt_id = kwargs.get('ppt_id', None)
                    if ppt_id is None:
                        ppt_id = GenerateID.create_random_id()
                        chat_root_id = GenerateID.create_random_id()
                        chat_id = GenerateID.create_random_id()

                        title = cls.get_title(content, model)
                        # 去除标题首位双引号
                        if title.startswith('"') and title.endswith('"'):
                            title = title[1:-1]

                        db.session.begin_nested()

                        ppt_base_model = PptBase(
                            ppt_id=ppt_id,
                            chat_root_id=chat_root_id,
                            user_id=user_id,
                            title=title,
                            outline=fetch_json(total_response),
                            status=0,
                            gen_type=0,
                        )
                        db.session.add(ppt_base_model)

                        ppt_chat_model = PptChat(
                            chat_parent_id=chat_root_id,
                            chat_id=chat_id,
                            content=kwargs.get('content'),
                            response=total_response,
                        )
                        db.session.add(ppt_chat_model)
                        db.session.commit()
                        yield f"event: done\ndata: {json.dumps({'ppt_id': ppt_id, 'title': title, 'chat_id': chat_id})}\n\n"
                    else:
                        ppt_base_model = db.session.query(PptBase).filter(PptBase.ppt_id == ppt_id)
                        ppt_base_info = ppt_base_model.first()

                        if ppt_base_info.title is None:
                            title = cls.get_title(content, model)
                            # 去除标题首位双引号
                            if title.startswith('"') and title.endswith('"'):
                                title = title[1:-1]
                            ppt_base_model.update({'title': title})

                        if ppt_base_info.chat_root_id is None:
                            chat_root_id = GenerateID.create_random_id()
                            ppt_base_model.update({'chat_root_id': chat_root_id})

                        ppt_base_model.update({'outline': fetch_json(total_response), 'status': 0, 'gen_type': 0})
                        chat_root_id = ppt_base_info.chat_root_id
                        chat_id = GenerateID.create_random_id()
                        ppt_chat_model = PptChat(
                            chat_parent_id=chat_root_id,
                            chat_id=chat_id,
                            content=kwargs.get('content'),
                            response=total_response,
                            # doc_location=kwargs.get('doc_location'),
                        )
                        db.session.add(ppt_chat_model)
                        db.session.commit()
                    yield f"event: done\ndata: {json.dumps({'ppt_id': ppt_id, 'chat_id': chat_id})}\n\n"
                except Exception as e:
                    db.session.rollback()
                    loggings.exception(1, e)
                    yield f"event: error\ndata: {json.dumps({'error': str(e)})}\n\n"

        headers = {
            'Content-Type': 'text/event-stream',
            'Cache-Control': 'no-cache',
            'Connection': 'keep-alive',
            'X-Accel-Buffering': 'no',
        }

        return Response(stream_with_context(generate()), mimetype="text/event-stream", headers=headers)

    @classmethod
    def get_prompt(cls, prompt, user_message: str) -> list:
        return [("system", prompt), ("user", user_message)]

    @classmethod
    def get_title(cls, user_input: str, model) -> list:
        if len(user_input) > 200:
            user_message = user_input[:200]
        prompt = ChatPromptTemplate.from_messages([("user", cls.title_prompt + '{user_input}')])
        chain = prompt | model
        title = chain.invoke({"user_input": user_input})
        return title.content

    @classmethod
    def get_chat_history(cls, ppt_id):
        """
            获取聊天记录
        """

        res_ppt_base = PptBaseController.get(ppt_id=ppt_id)
        if res_ppt_base['code'] != RET.OK or len(res_ppt_base['data']) != 1:
            return {'code': RET.DATAERR, 'message': 'PPT not found', 'data': []}
        ppt_base = res_ppt_base['data'][0]

        if ppt_base['chat_root_id'] is None:
            return {'code': RET.DATAERR, 'message': 'Chat history not found', 'data': []}
        res_ppt_chat = cls.get(chat_parent_id=ppt_base['chat_root_id'])
        if res_ppt_chat['code'] != RET.OK:
            return res_ppt_chat
        return {'code': RET.OK, 'message': 'OK', 'data': res_ppt_chat['data']}

    @classmethod
    def get_formatted_chat_history(cls, ppt_id):
        """
            格式化聊天记录
        """
        res_data = cls.get_chat_history(ppt_id=ppt_id)
        if res_data['code'] != RET.OK:
            return res_data
        data = []
        for chat in res_data['data']:
            """
                chat_parent_id=kwargs.get('chat_parent_id'),
                chat_id=kwargs.get('chat_id'),
                content=kwargs.get('content'),
                response=kwargs.get('response'),
            """
            data.append({
                'is_bot': False,
                'chat_id': chat['chat_id'],
                'message': chat['content'],
            })
            data.append({
                'is_bot': True,
                'chat_id': chat['chat_id'],
                'message': chat['response'],
            })
        return {'code': RET.OK, 'message': 'OK', 'data': data}
